from manimlib import *
import numpy as np


class convolution(Scene):
    def construct(self):
        x_range = (-5, 8)
        y_range = (-3, 9)
        axes = Axes(x_range=x_range, y_range=y_range)
        axes.add_coordinate_labels(range(x_range[0]+1,x_range[1],2),range(y_range[0]+1,y_range[1],2))
        inputs = [1,4,6,5,3,1]
        responses = []
        bottoms = []
        functions = []
        total_points = []
        sum_points = []
        input_points = VGroup()
        total_response = VGroup()
        add_points = VGroup()
        for index, item in enumerate(inputs):  # 在inputs定义域内的点的计算
            function = axes.get_graph(lambda x: item*np.exp(-x+index), x_range=(index,x_range[1])).set_color(YELLOW)
            bottom = axes.get_graph(lambda x: np.zeros(x.shape), x_range=(index,x_range[1])).set_color(GREY)
            point = Dot(axes.c2p(index,item)).set_color(BLUE)
            vline = axes.get_v_line_to_graph(index, function)

            response = function.copy()
            response.add_line_to(axes.c2p(x_range[1],0))
            response.add_line_to(axes.c2p(index,0))
            response.add_line_to(function.get_start())

            row_points = VGroup()
            row_sum = 0
            for i in range(index+1):
                y = inputs[i]*np.exp(-index+i)
                row_sum += y
                add_point = Dot(axes.c2p(index,y)).set_color(YELLOW)
                add_points.add(add_point)
                row_points.add(add_point)
            sum_point = Dot(axes.c2p(index,row_sum)).set_color(ORANGE)
                
            responses.append(response.set_fill(YELLOW,opacity=0.5))
            bottoms.append(bottom)
            functions.append(function)
            total_points.append(row_points)
            sum_points.append(sum_point)
            input_points.add(VGroup(point, vline))
            total_response.add(response)

        for pro in range(2):  # 在inputs定义域外延展两点
            index = len(inputs) + pro
            row_sum = 0
            row_points = VGroup()
            for i in range(len(inputs)):
                y = inputs[i]*np.exp(-index+i)
                row_sum += y
                add_point = Dot(axes.c2p(index,y)).set_color(YELLOW)
                add_points.add(add_point)
                row_points.add(add_point)
            sum_points.append(Dot(axes.c2p(index,row_sum)).set_color(ORANGE))
            total_points.append(row_points)
        
        impact_response = responses[0].copy()
        get_points = add_points.copy()
        text1 = Text('This is an \n\nimpact response.').scale(0.5).to_corner(UL)
        text2 = Text('This is an \n\ninput signal.').scale(0.5).to_corner(UL)
        text3 = Text('When impact response \n\nmeet input signal.').scale(0.5).to_corner(UL)
        text4 = Text('In the view of\n\ndiscrete time.').scale(0.5).to_corner(UL)
        text5 = Text('Add all the responses together,\n\nand we can get\n\nthe convolution.').scale(0.5).to_corner(UL)
        text6 = Text('What if in the view of\n\ncontinuous time?')

        self.play(FadeIn(axes))
        self.play(Write(text1))
        self.play(FadeIn(impact_response.set_color(RED)))
        self.wait(2)
        self.play(FadeOut(impact_response), Transform(text1,text2))
        self.play(FadeIn(input_points))
        self.wait(2)
        self.play(Uncreate(text1))
        self.play(Write(text3))
        for i in range(len(inputs)):
            self.play(Transform(bottoms[i],functions[i]), run_time=0.5)
            self.play(FadeIn(responses[i]), FadeOut(bottoms[i]), run_time=0.5)
        self.play(FadeOut(input_points), Uncreate(text3))
        self.play(Write(text4))
        self.play(FadeOut(total_response), FadeIn(add_points))
        self.play(Uncreate(text4))
        self.play(Write(text5))
        self.wait(1)
        for i in range(len(total_points)):
            self.play(Transform(total_points[i],sum_points[i]), run_time=0.7)
        self.wait(3)
        self.play(FadeOut(axes), FadeOut(add_points), FadeOut(text5))


        x_range = (-5, 8)
        y_range = (-3, 4)
        axes1 = Axes(x_range=x_range, y_range=y_range)
        axes1.add_coordinate_labels(range(x_range[0]+1,x_range[1],2),range(y_range[0]+1,y_range[1],2))

        bottoms[0] = axes1.get_graph(lambda x: np.zeros(x.shape), x_range=(0,x_range[1])).set_color(GREY)
        functions[0] = axes1.get_graph(lambda x: np.exp(-x), x_range=(0,x_range[1])).set_color(YELLOW)
        vline = axes1.get_v_line_to_graph(0, functions[0])
        impact_response = VGroup(functions[0],vline)
        
        responses = []
        total_response = VGroup()
        for i in range(x_range[1]):
            response = axes1.get_graph(lambda x: np.exp(-x+i), x_range=(i,x_range[1])).set_color(YELLOW).set_opacity(0.2)
            response.add_line_to(axes1.c2p(x_range[1],0))
            response.add_line_to(axes1.c2p(i,0))
            response.add_line_to(response.get_start())
            responses.append(response.set_fill(YELLOW, opacity=0.1))
            total_response.add(response)
        step = axes1.get_graph(lambda x: np.full(x.shape,1), x_range=(0,x_range[1])).set_color(BLUE)
        sum_response = axes1.get_graph(lambda x: np.full(x.shape,1)-np.exp(-x), x_range=(0,x_range[1])).set_color(ORANGE)
        
        text7 = Text('This is an\n\nunit step signal.').scale(0.5).to_corner(UL)
        text8 = Text('This is\n\nthe response.').scale(0.5).to_corner(UL)
        text9 = Text('To sum up,\n\nwe can get \n\nthe convolution.').scale(0.5).to_corner(UL)

        self.play(Write(text6))
        self.wait(2)
        self.play(FadeOut(text6), FadeIn(axes1))
        self.play(ShowCreation(text7), FadeIn(step))
        self.wait(2)
        self.play(Uncreate(text7), ShowCreation(text8))
        self.play(Transform(bottoms[0],functions[0]), run_time=0.5)
        self.play(FadeOut(bottoms[0]), FadeIn(impact_response))
        for i in range(x_range[1]):
            self.play(ApplyMethod(impact_response.shift,axes1.c2p(1,0)-axes1.c2p(0,0)), FadeIn(responses[i]), run_time=1, rate_func=linear)
        self.play(FadeOut(impact_response), FadeOut(step))
        self.play(Uncreate(text8))
        self.play(Write(text9))
        self.play(ApplyMethod(total_response.set_opacity, 0.5))
        self.play(Transform(bottoms[0], sum_response), FadeOut(total_response))
        self.wait(3)
        self.play(FadeOut(axes1), FadeOut(bottoms[0]), FadeOut(text9))


        RC = SVGMobject('RC_circuit.svg').set_fill(opacity=0)
        Us = Tex('U_s').scale(0.5).shift(LEFT*1)
        RC_circuit = VGroup(RC, Us)

        zi_axes = axes1.copy()
        zi_function = zi_axes.get_graph(lambda x: np.exp(-x), x_range=(0,x_range[1])).set_color(YELLOW)
        zi_label = Text('zero input response').scale(0.7).shift(DOWN*3.5)
        zi_response = VGroup(zi_axes, zi_function, zi_label).scale(0.5).to_corner(UR)

        zs_axes = axes1.copy()
        zs_function = zs_axes.get_graph(lambda x: np.full(x.shape,1)-np.exp(-x), x_range=(0,x_range[1])).set_color(ORANGE)
        zs_label = Text('zero state response').scale(0.7).shift(DOWN*3.5)
        zs_response = VGroup(zs_axes, zs_function, zs_label).scale(0.5).to_corner(DR)
        
        text10 = Tex('h(t)=e^{-t}u(t)').shift(LEFT*4+UP*3)
        text11 = Tex('x(t)=u(t)').shift(LEFT*4+UP*2)
        text12 = Tex(r'Y(s)')
        text12_1 = Tex(r'=X(s)\cdot H(s)')
        VGroup(text12[0], text12_1).arrange(RIGHT).shift(LEFT*4)
        text12_2 = Tex(r'=\frac{1}{s}\cdot\frac{1}{s+1}')
        VGroup(text12[0], text12_2).arrange(RIGHT).shift(LEFT*4)
        text12_3 = Tex(r'=\frac{1}{s}-\frac{1}{s+1}')
        VGroup(text12[0], text12_3).arrange(RIGHT).shift(LEFT*4)
        text13 = Tex(r'y(t)=u(t)-e^{-t}u(t)').shift(LEFT*4+DOWN)
        all_text = VGroup(text10, text11, text12, text12_3)
        text_bye = Text('Thanks for watching!')

        self.play(GrowFromCenter(RC_circuit.scale(2)))
        self.wait(2)
        self.play(ApplyFunction(lambda mob: mob.scale(0.5).shift(LEFT*3), RC_circuit))
        self.play(GrowFromEdge(zi_response, DOWN))
        self.wait(2)
        self.play(GrowFromEdge(zs_response, DOWN))
        self.wait(1)
        self.play(ShrinkToCenter(RC_circuit))
        self.play(Write(text10), Write(text11))
        self.wait(2)
        self.play(Write(text12), Write(text12_1))
        self.wait(1)
        self.play(FadeOut(text12_1), Write(text12_2))
        self.wait(1)
        self.play(FadeOut(text12_2), Write(text12_3))
        self.wait(1)
        self.play(Write(text13))
        self.wait(2)
        self.play(FadeOut(all_text), FadeOut(zi_response), ApplyMethod(text13.shift, RIGHT*6+UP*2.5), ApplyFunction(lambda mob: mob.move_to(np.array([0,0,0])).scale(2), zs_response))
        self.play(FadeOut(zs_label))
        self.wait(2)
        self.play(FadeOut(zs_axes), FadeOut(zs_function), FadeOut(text13))


        discrete_view = VGroup(axes, get_points)
        continuous_view = VGroup(axes1, total_response.set_opacity(0.2), step)

        self.play(FadeIn(discrete_view))
        self.play(ApplyFunction(lambda mob: mob.shift(LEFT*3).scale(0.4), discrete_view))
        self.wait(1)
        self.play(FadeIn(continuous_view))
        self.play(ApplyFunction(lambda mob: mob.shift(RIGHT*3).scale(0.4), continuous_view))
        self.wait(2)
        self.play(FadeOut(discrete_view), FadeOut(continuous_view))
        
        self.play(ShowCreation(text_bye))
